# -*- coding: utf-8 -*-

import pandas as pd
from jqdatapy.api import get_fundamentals
from pandas._libs.tslibs.timedeltas import Timedelta

from zvt.contract.api import df_to_db
from zvt.contract.recorder import TimeSeriesDataRecorder
from zvt.domain import Stock, StockValuation, Etf
from zvt.recorders.joinquant.common import to_jq_entity_id
from zvt.utils.time_utils import now_pd_timestamp, to_time_str, to_pd_timestamp
from zvt.recorders.em import em_api


class EMChinaStockValuationRecorder(TimeSeriesDataRecorder):
    entity_provider = "em"
    entity_schema = Stock
    # 数据来自em
    provider = "em"

    data_schema = StockValuation

    def record(self, entity, start, end, size, timestamps):
        if start is None:
            self.logger.error(
                "entity_code:{},entity_name:{}  cann't find start.可能已退市".format(
                    entity.code,entity.name
                )
            )
            return None
        start = max(start, to_pd_timestamp("2005-01-01"))
        end = min(now_pd_timestamp(), start + Timedelta(days=500))

        count: Timedelta = end - start
        df = em_api.get_stock_valuation(entity.code)

        df["entity_id"] = entity.id
        df["code"] = entity.code
        df["name"] = entity.name
        df["id"] = df["timestamp"].apply(lambda x: "{}_{}".format(entity.id, to_time_str(x)))
        df = df.rename(
            {"total_mv": "market_cap"},
            axis="columns",
        )

        df_to_db(df=df, data_schema=self.data_schema, provider=self.provider, force_update=self.force_update)

        return None


if __name__ == "__main__":
    # 上证50
    # df = Etf.get_stocks(code="510050")
    # stocks = df.stock_id.tolist()
    # print(stocks)
    # print(len(stocks))

    EMChinaStockValuationRecorder(day_data=True, codes=["000508"], force_update=False).run()

# the __all__ is generated
__all__ = ["EMChinaStockValuationRecorder"]
